from rwkv_cpp.rwkv_cpp_model import RWKVModel
from rwkv_cpp.rwkv_cpp_shared_library import load_rwkv_shared_library
from tokenizer_util import get_tokenizer
import sampling

from typing import List, Callable, Optional
from numpy import ndarray
import re

class RWKVModelStateExecutives:

    def __init__(self, rwkv: RWKVModel) -> None:
        self._state: Optional[ndarray] = None
        self._rwkv: RWKVModel = rwkv

    def evalSequence(self, tokens: List[int]) -> ndarray:
        logits, state = self._rwkv.eval_sequence_in_chunks(
            tokens, self._state, use_numpy=True
        )
        self._state = state
        return logits

    def eval(self, token: int) -> ndarray:
        logits, state = self._rwkv.eval(
            token, self._state, use_numpy=True
        )
        self._state = state
        return logits

    def reset(self) -> None:
        self._state = None

class QAStateExecutives:

    def __init__(
        self,
        se: RWKVModelStateExecutives,
        decoder: Callable[[List[int]], str],
        encoder: Callable[[str], List[int]]
    ) -> None:
        self._se: RWKVModelStateExecutives = se
        self._decoder: Callable[[List[int]], str] = decoder
        self._encoder: Callable[[str], List[int]] = encoder
        self._prefix: str = ''
        self._prefixTokens: List[int] = []

    def handle(self, prompt: str, callback: Callable[str, None]) -> None:
        fullOutText = ''
        tokens = self._prefixTokens + self._encoder(self._prefix + prompt)
        logits = self._se.evalSequence(tokens)

        while True:
            token: int = sampling.sample_logits(logits, top_p=0.5, temperature=0.8)
            if 0 == token: # next token is end of sequence
                self._prefix = '\n\n'
                self._prefixTokens = []
                return

            decodeString = self._decoder([token])
            fullOutText = fullOutText + decodeString
            if fullOutText.endswith('\n\n'):
                self._prefix = ''
                self._prefixTokens = [token]
                return
            callback(decodeString)

            logits = self._se.eval(token)

class Chat:

    def __init__(self, model: str) -> None:
        rwkv: RWKVModel = RWKVModel(load_rwkv_shared_library(), model, 7)
        decoder, encoder = get_tokenizer('auto', rwkv.n_vocab)
        self._se: QAStateExecutives = QAStateExecutives(
            RWKVModelStateExecutives(rwkv),
            decoder,
            encoder
        )
        self._lastPrint = ''

    def run(self) -> None:
        while True:
            question = input('Question: ')
            print('\nAnswer:', end='', flush=True)
            self.answer(question, self.print)
            if not self._lastPrint.endswith('\n'):
                print('\n', flush=True, end='')
            print('\n', flush=True, end='')

    def answer(self, question: str, callback: Callable[str, None]):
        self._se.handle(self.getPrompt(question), callback)

    def getPrompt(self, question: str) -> str:
        formatted = re.sub(r'[\r|\n]+', '\n', question).strip()
        prompt = f'Question: {formatted}\n\nAnswer:'
        return prompt

    def print(self, x) -> None:
        self._lastPrint = x
        print(x, flush=True, end='')

if __name__ == '__main__':
    Chat('D:\\my_files\\rwkv.cpp-7B-Q5_1.bin').run()
